import open3d as o3d
import numpy as np
import yaml


def load_graph_from_yaml(file_path):
    """从 YAML 文件加载图结构"""
    with open(file_path, "r") as f:
        graph = yaml.safe_load(f)
    nodes = np.array(graph["nodes"], dtype=np.float64)
    edges = np.array(graph["edges"], dtype=np.int32)
    return nodes, edges


def create_3d_graph_visualization(nodes, edges):
    """创建 3D 可视化图结构"""
    # 将 2D 节点坐标转换为 3D（z = 0）
    nodes_3d = nodes

    # 创建点云表示节点
    point_cloud = o3d.geometry.PointCloud()
    point_cloud.points = o3d.utility.Vector3dVector(nodes_3d)

    # 创建线集合表示边
    line_set = o3d.geometry.LineSet()
    line_set.points = o3d.utility.Vector3dVector(nodes_3d)
    line_set.lines = o3d.utility.Vector2iVector(edges)

    # 设置颜色
    point_cloud.paint_uniform_color([1.0, 0.0, 0.0])  # 红色节点
    line_set.paint_uniform_color([0.0, 0.0, 1.0])  # 蓝色边

    # 设置点的大小和线的宽度
    point_size = 5.0
    line_width = 2.0
    vis = o3d.visualization.Visualizer()
    vis.create_window(window_name="3D Graph Visualization", width=800, height=600)

    # 添加几何体
    vis.add_geometry(point_cloud)
    vis.add_geometry(line_set)

    # 设置渲染参数
    render_option = vis.get_render_option()
    render_option.point_size = point_size
    render_option.line_width = line_width
    render_option.background_color = np.array([0.1, 0.1, 0.1])  # 深色背景

    # 设置初始视角
    view_control = vis.get_view_control()
    view_control.set_front([0.0, 0.0, -1.0])  # 正视图
    view_control.set_up([0.0, 1.0, 0.0])  # 上方向
    view_control.set_zoom(0.5)

    return vis


def main(yaml_file):
    # 加载图结构
    print(f"[1/2] 正在加载图结构：{yaml_file}")
    try:
        nodes, edges = load_graph_from_yaml(yaml_file)
        print(f"  >> 节点数量: {len(nodes)}, 边数量: {len(edges)}")
    except Exception as e:
        print(f"  >> 加载失败: {e}")
        return

    # 创建可视化器
    print("[2/2] 正在创建 3D 可视化...")
    vis = create_3d_graph_visualization(nodes, edges)

    # 运行可视化器
    print("  >> 启动 3D 可视化窗口，请使用鼠标操作视角...")
    vis.run()
    vis.destroy_window()


if __name__ == "__main__":
    path = "../maps/multifloor_graph.yaml"
    main(path)
